package com.tong.lettuce.mutex;

import io.lettuce.core.RedisClient;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.sync.RedisCommands;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;

public class Lock {
        private Logger logger = LoggerFactory.getLogger(Lock.class);
        /**
         * 当前线程的锁集合
         */
        private ThreadLocal<Map<String, Integer>> lockers = new ThreadLocal<>();
        /**
         * 当前线程锁的key和value集合
         */
        private ThreadLocal<Map<String, String>> values = new ThreadLocal<>();
        private RedisClient client;
        private RedisCommands<String, String> commands;
        /**
         * 两个原子变量，用于存储加锁和解锁脚本的sha ID
         */
        private final AtomicReference<String> LUA_SHA_LOCK = new AtomicReference<>();
        private final AtomicReference<String> LUA_SHA_UNLOCK = new AtomicReference<>();

        /**
         * Lock的构造函数
         * @param client RedisClient的实例
         */
        public Lock(RedisClient client) {
            this.client = client;
            StatefulRedisConnection<String, String> connection = client.connect();
            commands = connection.sync();
            // 加锁的lua脚本
            // 键不存在则创建，并设置键的的过期时间为5s, 如果已存在则创建失败，
            final String lockLua = "local result = redis.call('setnx', KEYS[1], ARGV[2]); " +
                    "if result == 1 then redis.call('pexpire', KEYS[1], tonumber(ARGV[1])) " +
                    "return nil else return redis.call('pttl', KEYS[1]) end";
            // 将脚本缓存到服务器，并保存它的唯一ID到原子变量中
            LUA_SHA_LOCK.compareAndSet(null, commands.scriptLoad(lockLua));
            // 释放锁的lua脚本
            final String unlockLua = "local result = redis.call('get', KEYS[1]);" +
                    "if result == ARGV[1] then redis.call('del', KEYS[1]) " +
                    "return 1 else return nil end";
            // 将脚本缓存到服务器，并保存它的唯一ID到原子变量中
            LUA_SHA_UNLOCK.compareAndSet(null, commands.scriptLoad(unlockLua));
        }

        /**
         * 尝试获取锁
         * @param key 锁的键
         * @return 是否成功获取锁
         */
        private Boolean tryLock(String key) {
            String[] keys = new String[]{key};
            // 失效时间为5秒
            String ttl = "5000";
            String value = getValueByKey(key);
            // 如果没有设置过值
            if (value == null) {
                // 锁的值使用UUID生成随机ID以保证值的唯一性
                value = UUID.randomUUID().toString();
                // 将新生成的值放入集合中
                values.get().put(key, value);
            }
            String[] args = new String[]{ttl, value};
            // 如果创建键成功，则说明加锁成功
            Long result = commands.evalsha(LUA_SHA_LOCK.get(), ScriptOutputType.INTEGER, keys, args);
            if (result == null) {
                return true;
            }
            boolean isLock;
            // 一直阻塞，直到拿到锁
            while (true) {
                try {
                    // 这里的实现是有问题的
                    // Redisson实现的分布式锁是使用发布订阅实现的(Java的reentrantLock通过队列实现的)
                    // 从而可以及时通知其它线程去抢锁
                    Thread.sleep(5);
                    // 继续尝试获取锁
                    isLock = this.tryLock(key);
                    if (isLock) {
                        return true;
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }

        /**
         * 尝试释放锁
         * @param key 锁的键
         * @return 是否成功释放锁
         */
        private boolean tryRelease(String key) {
            String[] keys = new String[]{key};
            String[] args = new String[]{getValueByKey(key)};
            // 释放锁
            Long result = commands.evalsha(LUA_SHA_UNLOCK.get(), ScriptOutputType.INTEGER, keys, args);
            return result != null;
        }

        /**
         * 获取当前线程的锁
         * @param key 锁的键
         * @return 当前线程的锁和该锁的重入次数
         */
        private Integer getLockerCnt(String key) {
            // 获取当前线程的锁集合
            Map<String, Integer> map = lockers.get();
            // 如果集合不为空，返回key对应的值
            if (map != null) {
                return map.get(key);
            }
            lockers.set(new HashMap<>(4));
            return null;
        }

        /**
         * 获取锁对应的值
         * @param key 锁的键
         * @return 锁对应的值
         */
        private String getValueByKey(String key) {
            // 获取当前线程的锁和对应值的键值对集合
            Map<String, String> map = values.get();
            // 如果集合不为空，返回key对应的值
            if (map != null) {
                return map.get(key);
            }
            values.set(new HashMap<>(4));
            return null;
        }

        /**
         * 加可重入锁
         * @param key 锁的键
         * @return 是否成功
         */
        public boolean lock(String key){
            Integer refCnt = getLockerCnt(key);
            if (refCnt != null) {
                // 如果锁已持有，则锁的引用计数加1
                lockers.get().put(key, refCnt + 1);
                return true;
            }
            // 尝试加锁
            boolean ok = this.tryLock(key);
            // 如果加锁失败，则返回
            if (!ok) {
                return false;
            }
            // 加锁成功，引用计数设置为1
            lockers.get().put(key, 1);
            return true;
        }

        /**
         * 释放可重入锁
         * @param key 锁的键
         * @return 是否成功
         */
        public boolean unlock(String key) {
            Integer refCnt = getLockerCnt(key);
            // 当前未持有锁
            if (refCnt == null) {
                return false;
            }
            // 锁的引用数减1
            refCnt --;
            // 引用计数大于0，说明还持有锁
            if (refCnt > 0) {
                lockers.get().put(key, refCnt);
            } else {
                // 否则从锁集合中删除该键，并释放锁
                lockers.get().remove(key);
                return this.tryRelease(key);
            }
            return true;
        }
    }
